import os

import SimpleITK as sitk
import cv2
from PIL import Image
import scipy.ndimage as ndimage
import numpy as np
import albumentations as A
'''

BTCV的标签
背景：0
肝（liver）：6
右肾（right kidney）：2
左肾（left kidney）：3
脾（spleen）：1
'''

def getLabel(label_array):
    '''
    1：肝脏
    2：肿瘤
    '''
    # label_array[label_array==1] = 0
    # label_array[label_array==2]=1
    return label_array
def getMaxAndMinFromNii(data_dir=r"/home/liukai/AllData/ATLAS2023/train",save_dir=""):
    '''
    '''
    label_dir = os.path.join(data_dir,"labelsTr" )
    image_dir = os.path.join(data_dir,"imagesTr" )
    subject_list = os.listdir(image_dir)
    subject_list.sort()
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    p_list=[]#求前百分之几的数，每个nii筛选出的列表
    for index, nii_name in enumerate(subject_list):
        # ---读取图像数据----#
        percentage=0.98
        image_path = os.path.join(image_dir, nii_name)
        image_nii = sitk.ReadImage(image_path)
        image_array = sitk.GetArrayFromImage(image_nii)
        max_old = image_array.max()
        min_old=image_array.min()
        n_array=image_array.flatten()
        n_array=np.sort(n_array)#升序
        p_max=n_array[int(len(n_array)*percentage):].min()
        p_min=n_array[:int(len(n_array)*(1-percentage))].min()

        image_array[image_array>p_max]=p_max
        image_array[image_array <p_min] = p_min
        print(subject_list[index],max_old,image_array.max(),'   ',min_old,image_array.min())
        new_image_nii=sitk.GetImageFromArray(image_array)
        sitk.WriteImage(new_image_nii,os.path.join(save_dir,subject_list[index]))
def processAndSaveATLAS(data_dir=r"/home/liukai/AllData/ATLAS2023/train",
                   save_dir=r"/home/liukai/AllData/liverTumorForDomain"):
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    number = 0
    spacing_z=3#z轴的分辨率统一
    label_dir = os.path.join(data_dir,"labelsTr" )
    image_dir = os.path.join(data_dir,"imagesTr" )
    subject_list=os.listdir(image_dir)
    subject_list.sort()
    for index,nii_name in (enumerate(subject_list)):
        #---读取图像数据----#
        image_path = os.path.join(image_dir, nii_name)
        image_nii = sitk.ReadImage(image_path)
        image_array = sitk.GetArrayFromImage(image_nii)
        label_path = os.path.join(label_dir, nii_name.replace('im', 'lb'))
        label_nii = sitk.ReadImage(label_path)
        label_array = sitk.GetArrayFromImage(label_nii)
        # print("提取目标区域前的shape",label_array.shape, image_array.shape)
        # -------------------提取image和label的肝脏和肿瘤区域-----------#
        x1 = np.any(label_array == 2, axis=(1, 2))
        x2 = np.any(label_array == 1, axis=(1, 2))
        have_liver_and_tumor = np.logical_or(x1, x2)
        indexs = list(np.where(have_liver_and_tumor))[0]
        start_slice = max(0, indexs[0] - 2)
        end_slice = min(indexs[-1], indexs[-1] + 2) + 1
        image_array = image_array[start_slice:end_slice]
        label_array = label_array[start_slice:end_slice]
        # print("提取目标区域后的shape",label_array.shape, image_array.shape)
        # -------------------处理图像数据---------------------#
        last_shape=np.array([len(image_array),256,256])
        image_dimension_adjustment = 1 / (image_array.shape / last_shape)
        image_dimension_adjustment[0]=image_nii.GetSpacing()[-1] / spacing_z
        image_array = ndimage.zoom(image_array, image_dimension_adjustment, order=3)  # 双线性插值
        #---z-score--#
        image_array = (image_array - image_array.mean()) / image_array.std()
        #---数值放缩到[-1,1]之间，2 * ((array - m) / (l - (m))) - 1
        image_array=2*((image_array-image_array.min())/(image_array.max()-image_array.min()))-1
        print(image_array.max(),image_array.min())
        # -------------------处理标注数据---------------------#
        label_array=getLabel(label_array)
        # print(set(label_array.flatten()))
        label_dimension_adjustment = 1 / (label_array.shape / last_shape)
        label_dimension_adjustment[0] = label_nii.GetSpacing()[-1] / spacing_z
        label_array = ndimage.zoom(label_array, label_dimension_adjustment, order=0)  # 只会出现原来出现过的数值
        # print("调整spacing后的shape:",image_array.shape,label_array.shape)
        #--------拼接image和label的每个2D切片，[image, mask] 。并保存-----#
        for i in range(len(image_array)):
            if 1 in label_array[i]:
                number += 1
                # 拼接并保存
                # combine_array = np.array([image_array[i], label_array[i]])
                last_dir=os.path.join(save_dir,nii_name.split('.')[0])
                if not os.path.exists(last_dir):
                    os.makedirs(last_dir)
                save_name = os.path.join(last_dir, str(i+1) + '.npz')
                np.savez(save_name, arr_0=image_array[i],arr_1=label_array[i])#为了跟代码的dateset对应，参数名和dataset统一arr_0，arr_1
    print(number)
def getSlice(label_array,n_class=2):
    '''
    获取三个方向上包含目标的切片的起始位置
    label_array:标注array
    n_class：类别数,不包含背景
    45 74     142 369     70 379
    45 74     149 398     66 390
    295 476     157 407     49 361
    '''
    #12,01,02
    spand_slice_z=0
    spand_slice_x = 3
    spand_slice_y = 3

    contain_class=None
    for i in range(1,n_class+1):
        if i==1:
            contain_class = np.any(label_array == i, axis=(1, 2))
        else:
            contain_class=np.logical_or(np.any(label_array == i, axis=(1, 2)), contain_class)
    indexs_z = list(np.where(contain_class))[0]
    start_slice_z = max(0, indexs_z[0] - spand_slice_z)
    end_slice_z = min(label_array.shape[0], indexs_z[-1] + spand_slice_z) + 1

    for i in range(1, n_class + 1):
        if i == 1:
            contain_class = np.any(label_array == i, axis=(0, 2))
        else:
            contain_class = np.logical_or(np.any(label_array == i, axis=(0, 2)), contain_class)
    indexs_x = list(np.where(contain_class))[0]
    start_slice_x = max(0, indexs_x[0] - spand_slice_x)
    end_slice_x = min(label_array.shape[1], indexs_x[-1] + spand_slice_x) + 1

    for i in range(1, n_class + 1):
        if i == 1:
            contain_class = np.any(label_array == i, axis=(0, 1))
        else:
            contain_class = np.logical_or(np.any(label_array == i, axis=(0, 1)), contain_class)
    indexs_y = list(np.where(contain_class))[0]
    start_slice_y = max(0, indexs_y[0] - spand_slice_y)
    end_slice_y = min(label_array.shape[2], indexs_y[-1] + spand_slice_y) + 1
    return start_slice_z,end_slice_z,start_slice_x,end_slice_x,start_slice_y,end_slice_y

def image_aug(image_array):
    transform = A.OneOf(
        [
            # A.ShiftScaleRotate(scale_limit=0, rotate_limit=15, p=1, border_mode=cv2.BORDER_REPLICATE),
            A.RandomResizedCrop(height=256, width=256, scale=(0.8, 1.2), p=0.5)
            # A.GaussNoise(var_limit=(0, 0.1),p=1)
            # A.RandomBrightnessContrast(brightness_limit=(-0.5, 0.5), contrast_limit=(-0.5, 0.5), p=1)
            # A.GaussianBlur(sigma_limit=(0.5, 1.5), p=1)
        ])
    for i in range(len(image_array)):
        transforme_image = transform(image=image_array[i])
        image_array[i] = transforme_image['image']
    return image_array

def image_to_show_pixel(image_array):
    image_array=255*(image_array-image_array.min())/(image_array.max()-image_array.min())
    return image_array
def processAndSaveLITS(data_dir=r"/home/liukai/AllData/ATLAS2023/train",
                   save_dir=r"/home/liukai/AllData/liverTumorForDomain"):
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    number = 0
    spacing_z=2#z轴的分辨率统一
    label_dir = os.path.join(data_dir,"segmentation" )
    image_dir = os.path.join(data_dir,"volume" )
    subject_list=os.listdir(image_dir)
    subject_list.sort()
    #阈值截断，看别的论文LITS数据集一般设为[-200,250]
    up=250
    down=-200

    for index,nii_name in enumerate(subject_list[:3]):
        #---读取图像数据----#
        image_path = os.path.join(image_dir, nii_name)
        image_nii = sitk.ReadImage(image_path)
        image_array = sitk.GetArrayFromImage(image_nii)
        label_path = os.path.join(label_dir, nii_name.replace('volume', 'segmentation'))
        label_nii = sitk.ReadImage(label_path)
        label_array = sitk.GetArrayFromImage(label_nii)
        # print("提取目标区域前的shape",label_array.shape, image_array.shape)
        # -------------------提取image和label的肝脏和肿瘤区域-----------#
        start_slice_z, end_slice_z, start_slice_x, end_slice_x, start_slice_y, end_slice_y=getSlice(label_array)
        print(start_slice_z,end_slice_z,'   ',start_slice_x,end_slice_x,'   ',start_slice_y,end_slice_y)
        image_array = image_array[start_slice_z:end_slice_z,start_slice_x:end_slice_x,start_slice_y:end_slice_y]
        label_array = label_array[start_slice_z:end_slice_z,start_slice_x:end_slice_x,start_slice_y:end_slice_y]
        # print("提取目标区域后的shape",label_array.shape, image_array.shape)
        # -------------------处理图像数据---------------------#
        last_shape=np.array([len(image_array),256,256])
        image_dimension_adjustment = 1 / (image_array.shape / last_shape)
        image_dimension_adjustment[0]=image_nii.GetSpacing()[-1] / spacing_z
        image_array = ndimage.zoom(image_array, image_dimension_adjustment, order=3)  # 双线性插值
        image_array[image_array>up]=up
        image_array[image_array<down]=down
        #-------z-score和max-min归一化----#
        image_array = (image_array - image_array.mean()) / image_array.std()
        image_array = 2 * ((image_array - image_array.min()) / (image_array.max() - image_array.min())) - 1
        image_array_aug=image_aug(image_array.copy())
        print(image_array.max(),image_array_aug.max())
        image_array=image_to_show_pixel(image_array)
        image_array_aug = image_to_show_pixel(image_array_aug)
        # -------------------处理标注数据---------------------#
        label_array=getLabel(label_array)
        # print(set(label_array.flatten()))
        label_dimension_adjustment = 1 / (label_array.shape / last_shape)
        label_dimension_adjustment[0] = image_nii.GetSpacing()[-1] / spacing_z
        label_array = ndimage.zoom(label_array, label_dimension_adjustment, order=0)  # 只会出现原来出现过的数值

        image_nii=sitk.GetImageFromArray(image_array)
        label_nii=sitk.GetImageFromArray(label_array)
        image_aug_nii = sitk.GetImageFromArray(image_array_aug)
        sitk.WriteImage(image_nii,os.path.join(save_dir,"volume-{}.nii".format(index)))
        sitk.WriteImage(image_aug_nii, os.path.join(save_dir, "volume_aug-{}.nii".format(index)))
        sitk.WriteImage(label_nii, os.path.join(save_dir,"segmentation-{}.nii".format(index)))

    print(number)

if __name__ == '__main__':
    #gpu2
    # processAndSaveATLAS(data_dir=r"/home/liukai/AllData/ATLAS2023/train",
    #                save_dir=r"/home/liukai/AllData/liverTumorForDomain/ATLAS2023AfterProcess")
    processAndSaveLITS(data_dir=r"D:\AllData\LITS",
                        save_dir=r"D:\AllData\LITS\cutSlice")
    getMaxAndMinFromNii(data_dir=r"D:\AllData\atlas-train-dataset-1.0.1\atlas-train-dataset-1.0.1\train",
                        save_dir=r"D:\AllData\atlas-train-dataset-1.0.1\atlas-train-dataset-1.0.1\processMaxPexil")

